package tools

import (
	"fmt"
	"gitee.com/jjawesomejj/awesome-util/helper/commonhelper"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"time"
)

type SshCommand struct {
	IpAddress   string
	Password    string
	Port        int
	UserName    string
	Stdout      *StdWriterWithPid
	Stderr      *StdWriterWithPid
	TimeOutFunc func() (time int, echoWords string, callBack func())
}

func (command *SshCommand) RunShell(shell string) (*SshCommand, error) {
	isFinish := false
	timeOut := time.Second * 30
	if command.TimeOutFunc != nil {
		timeoutInt, _, _ := command.TimeOutFunc()
		timeOut = time.Duration(timeoutInt) * time.Second
	}
	client, err := sshClient(command.IpAddress, command.UserName, command.Port, timeOut)
	if err != nil {
		return command, err
	}

	session, err := client.NewSession()
	if err := session.RequestPty("xterm", 80, 40, ssh.TerminalModes{}); err != nil {
		return command, err
	}
	if err != nil {
		command.Stdout.Write([]byte(err.Error()))
		return command, err
	}
	defer func() {
		session.Close()
		isFinish = true
	}()
	command.Stdout.Write([]byte(shell + "\n"))
	session.Stdout = command.Stdout
	session.Stderr = command.Stderr
	if command.TimeOutFunc != nil {
		timeout, echoWords, callback := command.TimeOutFunc()
		time.AfterFunc(time.Duration(timeout)*time.Second, func() {
			if isFinish == false {
				command.Stdout.Write([]byte(commonhelper.Format(echoWords+",超时时间{}", strconv.Itoa(timeout))))
				session.Close()
				if callback != nil {
					callback()
				}
			}
		})
	}
	err = session.Run(shell)
	if err != nil {
		command.Stdout.Write([]byte(err.Error()))
		return command, err
	}
	return command, err
}

func sshClient(IpAddress string, username string, port int, timeout time.Duration) (*ssh.Client, error) {
	var (
		addr         string
		clientConfig *ssh.ClientConfig
		client       *ssh.Client
		err          error
	)

	homePath, err := os.UserHomeDir()
	if err != nil {
		return nil, err
	}
	key, err := ioutil.ReadFile(path.Join(homePath, ".ssh", "id_rsa"))
	if err != nil {
		return nil, err
	}
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, err
	}

	clientConfig = &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		Timeout:         timeout,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	// connet to ssh
	addr = fmt.Sprintf("%s:%d", IpAddress, port)

	if client, err = ssh.Dial("tcp", addr, clientConfig); err != nil {
		return nil, err
	}

	return client, nil
}
